package main

import (
	"context"
	"fmt"
	"log"
	"os"
	"strings"
	"time"

	"mcp-server/internal/qdrant"
	"mcp-server/internal/storage"
	"mcp-server/internal/vector"
)

// VectorVerifier 向量验证工具
type VectorVerifier struct {
	storageManager *storage.StorageManager
	vectorService  vector.VectorService
	ctx            context.Context
}

// NewVectorVerifier 创建向量验证工具
func NewVectorVerifier() (*VectorVerifier, error) {
	ctx := context.Background()

	// 初始化存储管理器
	storageConfig := &storage.StorageConfig{
		QdrantConfig: &qdrant.Config{
			DataPath: "./knowledge_vectors",
			Collection: qdrant.CollectionConfigYAML{
				Name:       "knowledge_base",
				VectorSize: 768,
				Distance:   "Cosine",
			},
		},
		BatchSize:   100,
		CacheSize:   1000,
		EnableCache: true,
	}
	storageManager := storage.NewStorageManager()
	if err := storageManager.Initialize(storageConfig); err != nil {
		return nil, fmt.Errorf("初始化存储管理器失败: %w", err)
	}

	// 初始化向量服务
	vectorService := vector.NewVectorService()
	serviceConfig := vector.DefaultServiceConfig()
	if err := vectorService.Initialize(serviceConfig); err != nil {
		return nil, fmt.Errorf("初始化向量服务失败: %w", err)
	}

	return &VectorVerifier{
		storageManager: storageManager,
		vectorService:  vectorService,
		ctx:            ctx,
	}, nil
}

// VerifyStorage 验证存储状态
func (vv *VectorVerifier) VerifyStorage() error {
	client := vv.storageManager.GetClient()

	fmt.Println("=== 存储状态验证 ===")

	// 检查集合是否存在
	exists, err := client.CollectionExists(vv.ctx, "knowledge_base")
	if err != nil {
		return fmt.Errorf("检查集合存在失败: %w", err)
	}
	fmt.Printf("集合 knowledge_base 存在: %v\n", exists)

	if !exists {
		fmt.Println("集合不存在，尝试创建...")
		collectionConfig := &storage.CollectionConfig{
			VectorSize: 768,
			Distance:   "cosine",
		}
		if err := client.CreateCollection(vv.ctx, "knowledge_base", collectionConfig); err != nil {
			return fmt.Errorf("创建集合失败: %w", err)
		}
		fmt.Println("集合创建成功")
	}

	// 获取集合信息
	info, err := client.GetCollectionInfo(vv.ctx, "knowledge_base")
	if err != nil {
		log.Printf("获取集合信息失败: %v", err)
	} else {
		fmt.Printf("集合信息:\n")
		fmt.Printf("  名称: %s\n", info.Name)
		fmt.Printf("  向量数量: %d\n", info.VectorCount)
		fmt.Printf("  向量维度: %d\n", info.VectorSize)
		fmt.Printf("  距离度量: %s\n", info.Distance)
		fmt.Printf("  状态: %s\n", info.Status)
	}

	// 获取存储统计
	stats, err := client.GetStorageStats(vv.ctx)
	if err != nil {
		log.Printf("获取存储统计失败: %v", err)
	} else {
		fmt.Printf("存储统计:\n")
		fmt.Printf("  总集合数: %d\n", stats.TotalCollections)
		fmt.Printf("  总向量数: %d\n", stats.TotalVectors)
	}

	return nil
}

// TestSearch 测试搜索功能
func (vv *VectorVerifier) TestSearch(queryText string) error {
	fmt.Printf("\n=== 搜索功能测试 ===\n")
	fmt.Printf("查询文本: %s\n", queryText)

	client := vv.storageManager.GetClient()

	// 创建测试向量（简化版本）
	testVector := make([]float32, 768)
	for i := range testVector {
		testVector[i] = 0.1 // 简单的测试向量
	}

	// 执行搜索
	results, err := client.SearchSimilar(vv.ctx, "knowledge_base", testVector, 5, 0.5)
	if err != nil {
		return fmt.Errorf("搜索失败: %w", err)
	}

	fmt.Printf("搜索结果:\n")
	fmt.Printf("  总匹配数: %d\n", results.Total)
	fmt.Printf("  查询时间: %v\n", results.QueryTime)
	fmt.Printf("  匹配结果:\n")

	for i, hit := range results.Hits {
		fmt.Printf("    %d. ID: %s, 相似度: %.4f\n", i+1, hit.ID, hit.Score)
		if hit.Payload != nil {
			if fileName, ok := hit.Payload["file_name"].(string); ok {
				fmt.Printf("       文件: %s\n", fileName)
			}
			if content, ok := hit.Payload["content"].(string); ok {
				// 截断内容显示
				if len(content) > 100 {
					content = content[:100] + "..."
				}
				fmt.Printf("       内容: %s\n", content)
			}
		}
	}

	return nil
}

// TestVectorization 测试向量化功能
func (vv *VectorVerifier) TestVectorization(testText string) error {
	fmt.Printf("\n=== 向量化功能测试 ===\n")
	fmt.Printf("测试文本: %s\n", testText)

	// 测试文本向量化
	result, err := vv.vectorService.VectorizeText(testText, map[string]interface{}{
		"test":      true,
		"timestamp": time.Now().Format(time.RFC3339),
	})
	if err != nil {
		return fmt.Errorf("向量化失败: %w", err)
	}

	fmt.Printf("向量化结果:\n")
	fmt.Printf("  成功: %v\n", result.Success)
	fmt.Printf("  向量维度: %d\n", len(result.Vector))
	fmt.Printf("  处理时间: %d ms\n", result.ProcessTime)

	if result.Quality != nil {
		fmt.Printf("  质量指标:\n")
		fmt.Printf("    向量范数: %.4f\n", result.Quality.VectorNorm)
		fmt.Printf("    文本长度: %d\n", result.Quality.TextLength)
		fmt.Printf("    一致性: %.4f\n", result.Quality.Consistency)
		fmt.Printf("    置信度: %.4f\n", result.Quality.Confidence)
	}

	// 显示向量的前10个维度
	fmt.Printf("  向量前10维: [")
	for i := 0; i < 10 && i < len(result.Vector); i++ {
		fmt.Printf("%.4f", result.Vector[i])
		if i < 9 && i < len(result.Vector)-1 {
			fmt.Printf(", ")
		}
	}
	fmt.Printf("...]\n")

	return nil
}

// RunInteractiveMode 运行交互模式
func (vv *VectorVerifier) RunInteractiveMode() {
	fmt.Println("\n=== 交互式验证模式 ===")
	fmt.Println("输入命令:")
	fmt.Println("  search <查询文本> - 搜索相似向量")
	fmt.Println("  vectorize <文本> - 测试向量化")
	fmt.Println("  stats - 显示统计信息")
	fmt.Println("  quit - 退出")

	for {
		fmt.Print("\n> ")
		var input string
		fmt.Scanln(&input)

		if input == "quit" {
			break
		}

		parts := strings.SplitN(input, " ", 2)
		command := parts[0]

		switch command {
		case "search":
			if len(parts) < 2 {
				fmt.Println("请提供查询文本")
				continue
			}
			if err := vv.TestSearch(parts[1]); err != nil {
				fmt.Printf("搜索失败: %v\n", err)
			}

		case "vectorize":
			if len(parts) < 2 {
				fmt.Println("请提供要向量化的文本")
				continue
			}
			if err := vv.TestVectorization(parts[1]); err != nil {
				fmt.Printf("向量化失败: %v\n", err)
			}

		case "stats":
			if err := vv.VerifyStorage(); err != nil {
				fmt.Printf("获取统计信息失败: %v\n", err)
			}

		default:
			fmt.Println("未知命令")
		}
	}
}

// Close 关闭资源
func (vv *VectorVerifier) Close() error {
	if vv.vectorService != nil {
		if err := vv.vectorService.Close(); err != nil {
			log.Printf("关闭向量服务失败: %v", err)
		}
	}

	if vv.storageManager != nil {
		if err := vv.storageManager.Close(); err != nil {
			log.Printf("关闭存储管理器失败: %v", err)
		}
	}

	return nil
}

func main() {
	// 创建验证工具
	verifier, err := NewVectorVerifier()
	if err != nil {
		log.Fatalf("创建验证工具失败: %v", err)
	}
	defer verifier.Close()

	// 验证存储状态
	if err := verifier.VerifyStorage(); err != nil {
		log.Fatalf("存储验证失败: %v", err)
	}

	// 测试向量化功能
	testText := "国资委负责监管国有企业的经营管理活动"
	if err := verifier.TestVectorization(testText); err != nil {
		log.Printf("向量化测试失败: %v", err)
	}

	// 测试搜索功能
	if err := verifier.TestSearch("国有企业监管"); err != nil {
		log.Printf("搜索测试失败: %v", err)
	}

	// 如果有命令行参数，运行交互模式
	if len(os.Args) > 1 && os.Args[1] == "-i" {
		verifier.RunInteractiveMode()
	}

	fmt.Println("\n验证完成！")
}
